#include <type_traits>
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;


template<typename TYPE_X1> class KernelSelectV2 {
    using T = TYPE_X1;
public:
    __aicore__ inline KernelSelectV2() {}
    __aicore__ inline void Init(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t totallength, uint32_t tileDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->totallength = totallength;
        this->tileDataNum = tileDataNum;

        conditionGm.SetGlobalBuffer((__gm__ int8_t*)condition, this->totallength + 32);
        x1Gm.SetGlobalBuffer((__gm__ T*)x1, this->totallength + 32);
        x2Gm.SetGlobalBuffer((__gm__ T*)x2, this->totallength + 32);
        yGm.SetGlobalBuffer((__gm__ T*)y, this->totallength + 32);

        pipe.InitBuffer(inQCondition, BUFFER_NUM, this->tileDataNum * sizeof(int8_t));
        pipe.InitBuffer(inQX1, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(inQX2, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(outQY, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp3, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(tmp5, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp6, this->tileDataNum * sizeof(half));
    }
    __aicore__ inline void Process() {

        int32_t loopCount = this->totallength / this->tileDataNum;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i,  this->tileDataNum);
            Compute(i, this->tileDataNum);
            CopyOut(i, this->tileDataNum);
        }
        if(this->totallength % this->tileDataNum != 0) {
            uint32_t taillength = this->totallength % this->tileDataNum;
            uint32_t taillengthalign = (taillength + 31) / 32 * 32;
            CopyIn(loopCount,  taillengthalign);
            Compute(loopCount, taillengthalign);
            CopyOut(loopCount, taillengthalign);
        }
    }
__aicore__ inline void InitE1(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t totallength, uint32_t tileDataNum, 
                                uint32_t x1_length, uint32_t x2_length, bool e1flag) {
        this->totallength = totallength;
        this->tileDataNum = tileDataNum;
        this->x1_length = x1_length;
        this->x2_length = x2_length;
        this->e1flag = e1flag;

        conditionGm.SetGlobalBuffer((__gm__ int8_t*)condition, this->totallength + 32);
        x1Gm.SetGlobalBuffer((__gm__ T*)x1, this->totallength + 32);
        x2Gm.SetGlobalBuffer((__gm__ T*)x2, this->totallength + 32);
        yGm.SetGlobalBuffer((__gm__ T*)y, this->totallength + 32);

        pipe.InitBuffer(inQBind, BUFFER_NUM, this->tileDataNum * sizeof(T));
     
    } 
    __aicore__ inline void ProcessE1() {
        if(this->e1flag == true) {
            int32_t loopCount = this->x1_length / this->tileDataNum;
            for (int32_t i = 0; i < loopCount; i++) {
                CopyInE1(i,  this->tileDataNum);
                CopyOutE1(i, this->tileDataNum);
            }
            if(this->x1_length % this->tileDataNum != 0) {
                uint32_t taillength = this->x1_length % this->tileDataNum;
                uint32_t taillengthalign = (taillength + 31) / 32 * 32;
                CopyInE1(loopCount,  taillengthalign);
                CopyOutE1(loopCount, taillengthalign);
            }            
        } else {
            int32_t loopCount = this->x2_length / this->tileDataNum;
            for (int32_t i = 0; i < loopCount; i++) {
                CopyInE1(i,  this->tileDataNum);
                CopyOutE1(i, this->tileDataNum);
            }
            if(this->x2_length % this->tileDataNum != 0) {
                uint32_t taillength = this->x2_length % this->tileDataNum;
                uint32_t taillengthalign = (taillength + 31) / 32 * 32;
                CopyInE1(loopCount,  taillengthalign);
                CopyOutE1(loopCount, taillengthalign);
            }        
        }
    }    
    __aicore__ inline void InitBroadcast(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t totallength, uint32_t tileDataNum,
                                uint32_t condition_length, uint32_t x1_length, uint32_t x2_length, 
                                int64_t numshapes, int64_t ss[], int64_t sf[]) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->totallength = totallength;
        this->tileDataNum = tileDataNum;
        this->numshapes = numshapes;
        this->condition_length = condition_length;
        this->x1_length = x1_length;
        this->x2_length = x2_length;
        for (int i = 0; i < 192; ++i) {
            ((int64_t *)this->shape)[i] = ss[i];
        }
        for(int i = 0; i < 64; ++i) {
            ((int64_t *)this->shapefull)[i] = sf[i];
        }        

        conditionGm.SetGlobalBuffer((__gm__ int8_t*)condition, this->totallength + 32);
        x1Gm.SetGlobalBuffer((__gm__ T*)x1, this->totallength + 32);
        x2Gm.SetGlobalBuffer((__gm__ T*)x2, this->totallength + 32);
        yGm.SetGlobalBuffer((__gm__ T*)y, this->totallength + 32);

        pipe.InitBuffer(inQCondition, BUFFER_NUM, this->tileDataNum * sizeof(int8_t));
        pipe.InitBuffer(inQX1, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(inQX2, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(outQY, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp3, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(tmp5, this->tileDataNum * sizeof(half));
        pipe.InitBuffer(tmp6, this->tileDataNum * sizeof(half));
    }
    __aicore__ inline void ProcessBroadcast() {

        int32_t loopCount = this->totallength / this->tileDataNum;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyInBR(i,  this->tileDataNum);
            Compute(i, this->tileDataNum);
            CopyOut(i, this->tileDataNum);
        }
        if(this->totallength % this->tileDataNum != 0) {
            uint32_t taillength = this->totallength % this->tileDataNum;
            uint32_t taillengthalign = (taillength + 31) / 32 * 32;
            CopyInBR(loopCount,  taillengthalign);
            Compute(loopCount, taillengthalign);
            CopyOut(loopCount, taillengthalign);
        }
    }                                       
private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length)
    {
        LocalTensor<int8_t> conditionLocal = inQCondition.AllocTensor<int8_t>();
        DataCopy(conditionLocal, conditionGm[progress * this->tileDataNum], length);
        inQCondition.EnQue(conditionLocal);

        LocalTensor<T> x1Local = inQX1.AllocTensor<T>();
        DataCopy(x1Local, x1Gm[progress * this->tileDataNum], length);
        inQX1.EnQue(x1Local);

        LocalTensor<T> x2Local = inQX2.AllocTensor<T>();
        DataCopy(x2Local, x2Gm[progress * this->tileDataNum], length);
        inQX2.EnQue(x2Local);
    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length)
    {
        LocalTensor<int8_t> conditionLocal = inQCondition.DeQue<int8_t>();
        LocalTensor<T> x1Local = inQX1.DeQue<T>();
        LocalTensor<T> x2Local = inQX2.DeQue<T>();
        LocalTensor<T> yLocal = outQY.AllocTensor<T>();
        if constexpr (std::is_same_v<T, int8_t> || std::is_same_v<T, signed char>) {
            auto conditionfp16 = tmp1.Get<half>();
            auto tmpfp16 = tmp2.Get<half>();
            auto tmpx1 = tmp5.Get<half>();
            auto tmpx2 = tmp6.Get<half>();
            Cast(conditionfp16, conditionLocal, RoundMode::CAST_NONE, length);
            Cast(tmpx1, x1Local, RoundMode::CAST_NONE, length);
            Cast(tmpx2, x2Local, RoundMode::CAST_NONE, length);
            Mul(tmpx1, tmpx1, conditionfp16, length);
            Duplicate(tmpfp16, half(1), length);
            Sub(tmpfp16, tmpfp16, conditionfp16, length);
            Mul(tmpx2, tmpx2, tmpfp16, length);
            Add(tmpx1, tmpx1, tmpx2, length);
            Cast(yLocal, tmpx1, RoundMode::CAST_TRUNC, length);
        } else if constexpr (std::is_same_v<T, int32_t>){
            auto conditionfp16 = tmp1.Get<half>();
            auto conditionint32 = tmp3.Get<int32_t>();
            Cast(conditionfp16, conditionLocal, RoundMode::CAST_NONE, length);
            Cast(conditionint32, conditionfp16, RoundMode::CAST_TRUNC, length);
            Mul(x1Local, x1Local, conditionint32, length);
            Duplicate(yLocal, int32_t(1), length);
            Sub(yLocal, yLocal, conditionint32, length);
            Mul(x2Local, x2Local, yLocal, length);
            Add(yLocal, x1Local, x2Local, length);

            // auto x1Localfp32 = tmp3.Get<float>();
            // auto x2Localfp32 = tmp4.Get<float>();
            // LocalTensor<half> y_fp16 = tmp1.Get<half>();

            // Cast(x1Localfp32, x1Local, RoundMode::CAST_ROUND, length);
            // Cast(x2Localfp32, x2Local, RoundMode::CAST_ROUND, length);

            // Sub(x1Localfp32, x2Localfp32, x1Localfp32, length);
            // Abs(x1Localfp32, x1Localfp32, length);     // x1 = ∣input−other∣

            // Abs(x2Localfp32, x2Localfp32, length);
            // Muls(x2Localfp32, x2Localfp32, this->rtol, length);
            // Adds(x2Localfp32, x2Localfp32, this->atol, length);     // x2 = atol+rtol×∣other∣

            // Max(x1Localfp32, x1Localfp32, x2Localfp32, length);
            // Sub(x1Localfp32, x2Localfp32, x1Localfp32, length);
            // Abs(x1Localfp32, x1Localfp32, length);
            // Mins(x1Localfp32, x1Localfp32, (float)MIN_ACCURACY_FP32, length);
            // Muls(x1Localfp32, x1Localfp32, (float)MAX_MUL_1_FP32, length);
            // Muls(x1Localfp32, x1Localfp32, (float)MAX_MUL_1_FP32, length);
            // Muls(x1Localfp32, x1Localfp32, (float)MAX_MUL_2_FP32, length);
            // Adds(x1Localfp32, x1Localfp32, (float)NEGATIVE_ONE_FP32, length);
            // Abs(x1Localfp32, x1Localfp32, length);

            // Cast(y_fp16, x1Localfp32, RoundMode::CAST_NONE, length);
            // Cast(yLocal, y_fp16, RoundMode::CAST_NONE, length);
        } else if constexpr (std::is_same_v<T, float>){

            // auto conditionfp16 = tmp1.Get<half>();
            // auto conditionfp32 = tmp3.Get<float>();
            // auto tmpfp32 = tmp4.Get<float>();
            // Cast(conditionfp16, conditionLocal, RoundMode::CAST_NONE, length);
            // Cast(conditionfp32, conditionfp16, RoundMode::CAST_NONE, length);
            // Mul(x1Local, x1Local, conditionfp32, length);
            // Duplicate(tmpfp32, float(1), length);
            // Sub(tmpfp32, tmpfp32, conditionfp32, length);
            // Mul(x2Local, x2Local, tmpfp32, length);
            // Add(yLocal, x1Local, x2Local, length);
            auto conditionfp16 = tmp1.Get<half>();
            auto conditionfp32 = tmp3.Get<float>();
            Cast(conditionfp16, conditionLocal, RoundMode::CAST_NONE, length);
            Cast(conditionfp32, conditionfp16, RoundMode::CAST_NONE, length);
            Mul(x1Local, x1Local, conditionfp32, length);
            Duplicate(yLocal, float(1), length);
            Sub(yLocal, yLocal, conditionfp32, length);
            Mul(x2Local, x2Local, yLocal, length);
            Add(yLocal, x1Local, x2Local, length);

        } else if constexpr (std::is_same_v<T, half>){

            auto conditionfp16 = tmp1.Get<half>();
            auto tmpfp16 = tmp2.Get<half>();
            Cast(conditionfp16, conditionLocal, RoundMode::CAST_NONE, length);
            Mul(x1Local, x1Local, conditionfp16, length);
            Duplicate(tmpfp16, half(1), length);
            Sub(tmpfp16, tmpfp16, conditionfp16, length);
            Mul(x2Local, x2Local, tmpfp16, length);
            Add(yLocal, x1Local, x2Local, length);

        }
        outQY.EnQue<T>(yLocal);
        inQCondition.FreeTensor(conditionLocal);
        inQX1.FreeTensor(x1Local);
        inQX2.FreeTensor(x2Local);
    }
    __aicore__ inline void CopyOut(int32_t progress, uint32_t length)
    {
        LocalTensor<T> yLocal = outQY.DeQue<T>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, length);
        outQY.FreeTensor(yLocal);
    }
    __aicore__ inline void CopyInE1(int32_t progress, uint32_t length)
    {
        LocalTensor<T> xLocal = inQBind.AllocTensor<T>();
        if(this->e1flag == true) {
            DataCopy(xLocal, x1Gm[progress * this->tileDataNum], length);
        } else {
            DataCopy(xLocal, x2Gm[progress * this->tileDataNum], length);
        }
        inQBind.EnQue(xLocal);
    }
    __aicore__ inline void CopyOutE1(int32_t progress, uint32_t length)
    {
        LocalTensor<T> xLocal = inQBind.DeQue<T>();
        DataCopy(yGm[progress * this->tileDataNum], xLocal, length);
        inQBind.FreeTensor(xLocal);
    } 
    __aicore__ inline void CopyInBR(int32_t progress, uint32_t length)
    {
        LocalTensor<int8_t> conditionLocal = inQCondition.AllocTensor<int8_t>();
        if(this->condition_length < this->totallength) {
            BroadCCondition(conditionLocal, progress * this->tileDataNum, length);
        } else {     
            DataCopy(conditionLocal, conditionGm[progress * this->tileDataNum], length);
        }    
        inQCondition.EnQue(conditionLocal);

        LocalTensor<T> x1Local = inQX1.AllocTensor<T>();
        if(this->x1_length < this->totallength) {
            BroadCX1(x1Local, progress * this->tileDataNum, length);
        } else {            
            DataCopy(x1Local, x1Gm[progress * this->tileDataNum], length);
        }    
        inQX1.EnQue(x1Local);

        LocalTensor<T> x2Local = inQX2.AllocTensor<T>();
        if(this->x2_length < this->totallength) {
            BroadCX2(x2Local, progress * this->tileDataNum, length);
        } else { 
            DataCopy(x2Local, x2Gm[progress * this->tileDataNum], length);
        }    
        inQX2.EnQue(x2Local);
    } 
    __aicore__ inline void BroadCCondition(LocalTensor<int8_t> &dst, uint32_t offset, uint32_t length) {

        for(uint32_t i = 0; i < length; i++) {
            int idxtmp = 0;
            int istart = i + offset;
            for(int k = 1; k <= this->numshapes; k++) {
                int kpos = 0;
                int krange = 1;
                if(k < this->numshapes) {
                    for(int m = k + 1; m <= this->numshapes; m++) {
                        krange *= shapefull[m - 1];
                    }
                    kpos = istart / krange;
                    istart = istart % krange;
                } else {
                    krange = shapefull[k - 1];
                    kpos = istart % krange;
                }
                //idxtmp += kpos * this->stride[k - 1];
                int krangeB = 1;
                if(shapefull[k - 1] == shape[0][k - 1]) {
                    if(k < this->numshapes) {
                        for(int m = k + 1; m <= this->numshapes; m++) {
                            krangeB *= shape[0][m - 1];
                        }
                        idxtmp += kpos * krangeB;
                    }  else {
                        idxtmp += kpos;
                    }
                }
            }
            T tmp = conditionGm.GetValue(idxtmp);
            dst.SetValue(i, tmp);
        }
    }    
    __aicore__ inline void BroadCX1(LocalTensor<T> &dst, uint32_t offset, uint32_t length) {
        if(this->x1_length == 1) {
            if constexpr (std::is_same_v<T, int32_t> || std::is_same_v<T, half> || std::is_same_v<T, float>) {
                T tmp = x1Gm.GetValue(0);
                Duplicate(dst, tmp, length);
                return;
            }
        }
        for(uint32_t i = 0; i < length; i++) {
            int idxtmp = 0;
            int istart = i + offset;
            for(int k = 1; k <= this->numshapes; k++) {
                int kpos = 0;
                int krange = 1;
                if(k < this->numshapes) {
                    for(int m = k + 1; m <= this->numshapes; m++) {
                        krange *= shapefull[m - 1];
                    }
                    kpos = istart / krange;
                    istart = istart % krange;
                } else {
                    krange = shapefull[k - 1];
                    kpos = istart % krange;
                }
                //idxtmp += kpos * this->stride[k - 1];
                int krangeB = 1;
                if(shapefull[k - 1] == shape[1][k - 1]) {
                    if(k < this->numshapes) {
                        for(int m = k + 1; m <= this->numshapes; m++) {
                            krangeB *= shape[1][m - 1];
                        }
                        idxtmp += kpos * krangeB;
                    }  else {
                        idxtmp += kpos;
                    }
                }
            }
            T tmp = x1Gm.GetValue(idxtmp);
            dst.SetValue(i, tmp);
        }
    }
    __aicore__ inline void BroadCX2(LocalTensor<T> &dst, uint32_t offset, uint32_t length) {
        if(this->x2_length == 1) {
            if constexpr (std::is_same_v<T, int32_t> || std::is_same_v<T, half> || std::is_same_v<T, float>) {
                T tmp = x2Gm.GetValue(0);
                Duplicate(dst, tmp, length);
                return;
            }    
        }
        for(uint32_t i = 0; i < length; i++) {
            int idxtmp = 0;
            int istart = i + offset;
            for(int k = 1; k <= this->numshapes; k++) {
                int kpos = 0;
                int krange = 1;
                if(k < this->numshapes) {
                    for(int m = k + 1; m <= this->numshapes; m++) {
                        krange *= shapefull[m - 1];
                    }
                    kpos = istart / krange;
                    istart = istart % krange;
                } else {
                    krange = shapefull[k - 1];
                    kpos = istart % krange;
                }
                //idxtmp += kpos * this->stride[k - 1];
                int krangeB = 1;
                if(shapefull[k - 1] == shape[2][k - 1]) {
                    if(k < this->numshapes) {
                        for(int m = k + 1; m <= this->numshapes; m++) {
                            krangeB *= shape[2][m - 1];
                        }
                        idxtmp += kpos * krangeB;
                    }  else {
                        idxtmp += kpos;
                    }
                }
            }
            T tmp = x2Gm.GetValue(idxtmp);
            dst.SetValue(i, tmp);
        }
    }          
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQCondition, inQX1, inQX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQY;
    TQueBind<TPosition::VECIN, TPosition::VECOUT, BUFFER_NUM> inQBind;
    TBuf<QuePosition::VECCALC> tmp1, tmp2, tmp3, tmp5, tmp6;
    GlobalTensor<int8_t> conditionGm;
    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;

    uint32_t totallength;
    uint32_t tileDataNum;
    bool e1flag;
    uint32_t condition_length;
    uint32_t x1_length;
    uint32_t x2_length;
    uint32_t alignnum;
    int64_t shape[3][64];
    int64_t numshapes;
    int64_t shapefull[192];
};

extern "C" __global__ __aicore__ void select_v2(GM_ADDR condition, GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelSelectV2<DTYPE_X1> op;
    if (TILING_KEY_IS(300)) {
        op.Init(condition, x1, x2, y, 
                tiling_data.totallength, tiling_data.tileDataNum);  
        op.Process();
    } else if (TILING_KEY_IS(200)) {
        op.InitBroadcast(condition, x1, x2, y, 
                tiling_data.totallength, tiling_data.tileDataNum,
                tiling_data.condition_length, tiling_data.x1_length, tiling_data.x2_length, 
                tiling_data.numshapes, tiling_data.shape, tiling_data.shapefull);  
        op.ProcessBroadcast();
    } else if (TILING_KEY_IS(100)) { 
        op.InitE1(condition, x1, x2, y, 
                tiling_data.totallength, tiling_data.tileDataNum,
                tiling_data.x1_length, tiling_data.x2_length,
                tiling_data.e1flag);  
        op.ProcessE1();
    }      
}